package com.wuwangfu.func;

import com.alibaba.druid.pool.DruidDataSource;
import org.apache.flink.api.java.tuple.Tuple2;
import org.apache.flink.configuration.Configuration;
import org.apache.flink.streaming.api.functions.async.ResultFuture;
import org.apache.flink.streaming.api.functions.async.RichAsyncFunction;

import java.sql.Connection;
import java.sql.PreparedStatement;
import java.sql.ResultSet;
import java.util.Collections;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.Future;
import java.util.function.Supplier;

/**
 * @Description：
 * @Author：jcshen
 * @Date：2023-06-29
 */
public class AsyncMySQLFunc extends RichAsyncFunction<String, Tuple2<String,String>> {
    //使用druid数据库连接池
    private transient DruidDataSource dataSource;
    //用于提交多个异步请求的线程池
    private transient ExecutorService executorService;
    //线程池最大线程数量
    private int maxConnTotal;

    public AsyncMySQLFunc(int maxConnTotal) {
        this.maxConnTotal = maxConnTotal;
    }

    @Override
    public void open(Configuration parameters) throws Exception {
        //创建固定大小的线程池
        executorService = Executors.newFixedThreadPool(maxConnTotal);
        //创建数据库连接池并指定对应的参数
        dataSource = new DruidDataSource();
        dataSource.setDriverClassName("com.mysql.jdbc.Driver");
        dataSource.setUsername("root");
        dataSource.setPassword("Changxin*8");
        dataSource.setUrl("jdbc://mysql:hadoop04:3306/bigdata?characterEncoding=UTF-8");
        dataSource.setMaxActive(maxConnTotal);
    }

    @Override
    public void asyncInvoke(String id, ResultFuture<Tuple2<String, String>> resultFuture) throws Exception {
        //调用线程池的submit方法，将查询请求丢入到线程池中异步执行，返回Future对象
        Future<String> future = executorService.submit(() -> {
            //查询数据库
            return queryFromMysql(id);
        });
        CompletableFuture.supplyAsync(new Supplier<String>() {
            @Override
            public String get() {
                try {
                    //获取查询结果
                    return future.get();
                } catch (Exception e) {
                    return null;
                }
            }
        }).thenAccept((String result)->{
            //将id和查询的结果用Tuple2封装，放入到 ResultFuture中输出
            resultFuture.complete(Collections.singleton(Tuple2.of(id,result)));
        });
    }

    private String queryFromMysql(String param) throws Exception {
        String sql = "select id,info from t_data where id = ?";
        String result = null;
        Connection connection = null;
        PreparedStatement stmt = null;
        ResultSet res = null;

        try {
            connection = dataSource.getConnection();
            stmt = connection.prepareStatement(sql);
            //设置查询参数
            stmt.setString(1,param);
            //执行查询
            res = stmt.executeQuery();
            while(res.next()){
                //返回查询结果
                result = res.getString("info");
            }
        } finally {
            if (res != null){
                res.close();
            }
            if (stmt != null){
                stmt.close();
            }
            if (connection !=null){
                connection.close();
            }
        }
        return  result;
    }

    @Override
    public void close() throws Exception {
        //关闭数据库连接池
        dataSource.close();
        //关闭线程池
        executorService.shutdown();
    }
}
